import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
import mnist_forward
import os

BATCH_SIZE = 200
LEARNING_RATE_BASE = 0.1
LEARNING_RATE_DECAY = 0.99
REGULARIZER = 0.0001
STEPS = 50000

MOVING_AVERAGE_DECAY = 0.99
MODEL_SAVE_PATH = "./model/"
MODEL_NAME = "mnist_model"


def backward(mnist):

    x = tf.placeholder(tf.float32, shape = (None, mnist_forward.INPUT_NODE))
    y_ = tf.placeholder(tf.float32, shape = (None,  mnist_forward.OUTPUT_NODE))

    y = mnist_forward.forward(x, REGULARIZER)

    global_step = tf.Variable(0,trainable = False)

    ce = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=y, labels=tf.argmax(y_,1))
    cem = tf.reduce_mean(ce)
    loss = cem + tf.add_n(tf.get_collection('losses'))


    '''
    # lossy与y_的差距，以MSE为例
    loss_mse = tf.reduce_mean(tf.square(y-y_))
    # 也可以是交叉熵和softmax的协同使用
    ce = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=y,labels=tf.argmax(y_,1))
    cem = tf.reduce_mean(ce)
    # 加入正则化损失——提高泛化性
    loss_total = loss_mse + tf.add_n(tf.get_collection('losses'))
    '''

    # 指数衰减学习率——加快优化的效率
    learning_rate = tf.train.exponential_decay(
        LEARNING_RATE_BASE, # 学习率基数，初始值
        global_step, # 几轮，计数器
        mnist.train.num_examples /BATCH_SIZE, # LEARNING_RATE_STEP, # 多少轮更新一次
        LEARNING_RATE_DECAY, # 衰减率
        staircase=True # False 学习率为平滑曲线 True 直线
    )
    
    # 定义训练过程
    train_step = tf.train.GradientDescentOptimizer(learning_rate).minimize(
        loss,global_step=global_step)

    # 滑动平均
    ema = tf.train.ExponentialMovingAverage(MOVING_AVERAGE_DECAY, global_step)
    # 更新列表 trainable_variables自动将所有带训练数据汇总为列表
    ema_op = ema.apply(tf.trainable_variables())
    with tf.control_dependencies([train_step,ema_op]):
        train_op = tf.no_op(name='train')
    
    # 实例化存储
    saver = tf.train.Saver()

    with tf.Session() as sess:
        init_op = tf.global_variables_initializer()
        sess.run(init_op)

        for i in range(STEPS):
            xs, ys = mnist.train.next_batch(BATCH_SIZE)
            _, loss_value, step = sess.run([train_op, loss, global_step], 
            feed_dict = {x:xs, y_:ys})
            # 每几轮打印出信息
            if i% 1000 == 0:
                print ("在 %d 步训练之后，训练集损失达到 %g" % (step, loss_value))
                saver.save(sess, os.path.join(MODEL_SAVE_PATH, MODEL_NAME), global_step=global_step)
        
        '''
        # 生成 xx yy 坐标网格
        xx,yy = np.mgrid[-3:3:.01, -3:3:.01]
        # 将xx，yy拉直，并合并成一个2列的矩阵，得到网格坐标点集合
        grid = np.c_[xx.ravel(),yy.ravel()]
        probs = sess.run(y, feed_dict = {x:grid})
        probs = probs.reshape(xx.shape)

    plt.scatter(X[:,0], X[:,1], c=np.squeeze(Y_c))
    plt.contour(xx,yy,probs,levels=[.5])  # x坐标轴，y坐标轴，该点高度，等高线高度
    plt.show()
'''

def main():
    mnist = input_data.read_data_sets('E:/1-CODE/database/MNIST_data/', one_hot=True)
    backward(mnist)

if __name__=='__main__': # 判断python运行的文件是否是主文件
    main()

    
